from pyspark.sql import SparkSession
from pyspark.sql.functions import col, sqrt, pow, max, round, avg, lit

from SparkSessionBase import SparkSessionBase  # 假设这是你自己的封装类


# 继承 SparkSessionBase 类
class DistanceApi(SparkSessionBase):
    SPARK_URL = "yarn"
    SPARK_APP_NAME = 'DistanceApi'
    ENABLE_HIVE_SUPPORT = True

    def __init__(self):
        self.spark = self._create_spark_session()
        self.spark.sparkContext.setLogLevel("ERROR")  # 设置日志级别为 ERROR，减少日志

    def run(self, longitude, latitude, status):
        business_df = self.spark.table('business')
        review_df = self.spark.table('review')  # 加入 review 表
        business_review_df = business_df.join(review_df, business_df.business_id == review_df.rev_business_id, "inner")

        # 计算 distance_km
        business_df_with_distance = business_review_df.withColumn(
            "distance_km",
            sqrt(
                pow(col("longitude") - lit(longitude), 2) +
                pow(col("latitude") - lit(latitude), 2)
            ) * lit(6371)
        )

        # 筛选和排序
        filtered_df = business_df_with_distance.filter(
            (col("distance_km") <= lit(1))
        ).orderBy(
            col("distance_km").asc()
        )

        # 选择需要的列
        final_df = filtered_df.select(
            "business_id",
            "name",
            "address",
            "categories",
            "city",
            "state",
            "postal_code",
            "stars",
            "review_count",
            "is_open",
            "distance_km",
            "attributes"  # 添加 attributes 列
        )

        # 对 distance_km 列进行四舍五入到两位小数
        final_df_with_rounded_distance = final_df.withColumn(
            "distance_km",
            round(col("distance_km"), 2)
        )

        # 根据 status 参数调整排序逻辑
        if status == 1:
            sorted_df = final_df_with_rounded_distance.orderBy(col("distance_km").asc())
        elif status == 2:
            sorted_df = final_df_with_rounded_distance.orderBy(col("stars").desc())
        elif status == 3:
            sorted_df = final_df_with_rounded_distance.orderBy(col("review_count").desc())
        else:
            raise ValueError("Invalid status value. Please use 1, 2, or 3.")

        # # 获取查询出的总数
        # total_count = sorted_df.count()
        # print(f"查询出的总数: {total_count}")
        #
        # # 显示结果
        # sorted_df.show(truncate=False)
        print("distance success")
        return sorted_df


if __name__ == '__main__':
    # 示例调用
    job = DistanceApi()
    job.run(
        longitude=-110.95723,
        latitude=32.28918,
        status=1
    )